
#include <sys/mman.h>
#include <fcntl.h>
#include <unistd.h>     // for getpid(), close()
#include <sys/types.h>  // for ftruncate()
#include <chrono>
#include <vector>
#include <string>
#include <thread>
#include <numa.h>
#include <numaif.h>
#include <cstring>
#include <stdio.h>
#include <omp.h>
#include <pthread.h>
#include <future>
#include <mutex>
#include <condition_variable>

#define ARRAY_SIZE (1024 * 1024 * 1024) // 1GB elements
#define NUM_TIMES 10
#define NUM_NUMA_NODES 4

// Structure to hold shared memory information
typedef struct {
    void* mapped_region;
    char *a, *b, *c;
    int fd;
    const char* name;
} SharedMemory;

void stream_copy(char* a, char* b, size_t size) {
    #pragma omp parallel for
    for (size_t i = 0; i < size; i++) {
        b[i] = a[i];
    }
}

void stream_copy1(char* a, char* b, size_t size, int target_numa) {
#pragma omp parallel
    {
        // Set NUMA affinity for each thread
        int thread_id = omp_get_thread_num();
        cpu_set_t cpuset;
        CPU_ZERO(&cpuset);
        
        // Get CPUs for this NUMA node
        for (int cpu = 0; cpu < numa_num_configured_cpus(); cpu++) {
            if (numa_node_of_cpu(cpu) == target_numa) {
                CPU_SET(cpu, &cpuset);
            }
        }
        
        // Set thread affinity
        pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset);
        
        #pragma omp for
        for(int j = 0; j < size; j++) {
            b[j] = a[j];
        }
    }
}

void stream_scale(char* a, char* b, char scalar, size_t size) {
    #pragma omp parallel for
    for (size_t i = 0; i < size; i++) {
        b[i] = scalar * a[i];
    }
}

void stream_add(char* a, char* b, char* c, size_t size) {
    #pragma omp parallel for
    for (size_t i = 0; i < size; i++) {
        c[i] = a[i] + b[i];
    }
}

void stream_triad(char* a, char* b, char* c, char scalar, size_t size) {
    #pragma omp parallel for
    for (size_t i = 0; i < size; i++) {
        a[i] = b[i] + scalar * c[i];
    }
}

int main() {
    // Initialize NUMA
    if (numa_available() < 0) {
        printf("NUMA is not available\n");
        return 1;
    }

    // Bind the main thread to NUMA node 0
    struct bitmask *cpu_mask = numa_allocate_cpumask();
    numa_bitmask_setbit(cpu_mask, 0);
    numa_bind(cpu_mask);
    numa_free_cpumask(cpu_mask);

    SharedMemory shm[NUM_NUMA_NODES];
    size_t shm_size = ARRAY_SIZE * sizeof(char) * 3; // For a, b, c arrays

    // Create and initialize shared memory regions on each NUMA node
    for (int node = 0; node < NUM_NUMA_NODES; node++) {
        char shm_name[32];
        snprintf(shm_name, sizeof(shm_name), "shm_pool1_%d", node);
        shm[node].name = strdup(shm_name);

        // Create shared memory
        shm[node].fd = shm_open(shm[node].name, O_CREAT | O_RDWR, 0666);
        if (shm[node].fd == -1) {
            perror("shm_open failed");
            return 1;
        }

        // Set size
        if (ftruncate(shm[node].fd, shm_size) == -1) {
            perror("ftruncate failed");
            return 1;
        }

        // Map memory
        shm[node].mapped_region = mmap(NULL, shm_size, PROT_READ | PROT_WRITE, MAP_SHARED, shm[node].fd, 0);
        if (shm[node].mapped_region == MAP_FAILED) {
            perror("mmap failed");
            return 1;
        }

       // // Bind the mapped region to specific NUMA node
        //struct bitmask *nodemask = numa_allocate_nodemask();
        //numa_bitmask_setbit(nodemask, node);
        numa_tonode_memory(shm[node].mapped_region, shm_size, node);
        //numa_free_nodemask(nodemask);

        // Setup arrays
        shm[node].a = (char*)shm[node].mapped_region;
        shm[node].b = shm[node].a + ARRAY_SIZE;
        shm[node].c = shm[node].b + ARRAY_SIZE;

        // Initialize data
        #pragma omp parallel for
        for (size_t i = 0; i < ARRAY_SIZE; i++) {
            shm[node].a[i] = 1;
            shm[node].b[i] = 2;
            shm[node].c[i] = 0;
        }
    }

    auto start = std::chrono::high_resolution_clock::now();
    for(int i = 0; i < 10; i++) {
        //stream_copy(static_cast<char*>(shm[0].mapped_region), 
        //static_cast<char*>(shm[0].mapped_region) + 512*1024*1024, 512*1024*1024);

        stream_copy(shm[0].a, shm[0].c, ARRAY_SIZE);

        
    }
    auto end = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
    printf("Rank 0: Time taken for 10 iterations: %d ms, bandwidth: %f GB/s\n", duration, 
    ((double)1024)*1024*1024*10/1e9/(duration/1000.0));
                    
    char scalar = 3;
    double times[NUM_NUMA_NODES][4][NUM_TIMES];
    
    printf("Running STREAM benchmark from NUMA node 0 to all NUMA nodes...\n");
    
    // Perform benchmark for each NUMA node's memory
    for (int node = 0; node < NUM_NUMA_NODES; node++) {
        printf("\nTesting NUMA node %d:\n", node);
        
        for (int t = 0; t < NUM_TIMES; t++) {
            struct timespec start, end;
            
            // Copy
            clock_gettime(CLOCK_MONOTONIC, &start);
            stream_copy(shm[node].a, shm[node].c, ARRAY_SIZE);
            clock_gettime(CLOCK_MONOTONIC, &end);
            times[node][0][t] = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9;
            
            // Scale
            clock_gettime(CLOCK_MONOTONIC, &start);
            stream_scale(shm[node].b, shm[node].c, scalar, ARRAY_SIZE);
            clock_gettime(CLOCK_MONOTONIC, &end);
            times[node][1][t] = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9;
            
            // Add
            clock_gettime(CLOCK_MONOTONIC, &start);
            stream_add(shm[node].a, shm[node].b, shm[node].c, ARRAY_SIZE);
            clock_gettime(CLOCK_MONOTONIC, &end);
            times[node][2][t] = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9;
            
            // Triad
            clock_gettime(CLOCK_MONOTONIC, &start);
            stream_triad(shm[node].a, shm[node].b, shm[node].c, scalar, ARRAY_SIZE);
            clock_gettime(CLOCK_MONOTONIC, &end);
            times[node][3][t] = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9;
        }

        // Calculate and print results for this NUMA node
        const char* labels[] = {"Copy", "Scale", "Add", "Triad"};
        for (int i = 0; i < 4; i++) {
            double min_time = times[node][i][0];
            for (int t = 1; t < NUM_TIMES; t++) {
                if (times[node][i][t] < min_time) min_time = times[node][i][t];
            }
            
            double bandwidth = 0.0;
            switch(i) {
                case 0: // Copy
                case 1: // Scale
                    bandwidth = 2 * sizeof(char) * ARRAY_SIZE / min_time / 1e9;
                    break;
                case 2: // Add
                case 3: // Triad
                    bandwidth = 3 * sizeof(char) * ARRAY_SIZE / min_time / 1e9;
                    break;
            }
            
            printf("%s: Best time = %f s, Bandwidth = %f GB/s\n", 
                   labels[i], min_time, bandwidth);
        }
    }

    // Cleanup
    for (int node = 0; node < NUM_NUMA_NODES; node++) {
        munmap(shm[node].mapped_region, shm_size);
        close(shm[node].fd);
        shm_unlink(shm[node].name);
        free((void*)shm[node].name);
    }
    
    return 0;
}
